from typing import List
from langchain.schema.embeddings import Embeddings
from common.utils import send_http


class MyEmbedding(Embeddings):

    def __init__(self, maas_model):
        self.api = maas_model

    def embed_documents(self, texts: List[str]) -> List[List[float]]:
        results = []
        for text in texts:
            result = self.api.embedings(text)
            results.append(result)
        return results

    def embed_query(self, text: str) -> List[float]:
        return self.api.embedings(text)


class MaasModel():

    def __init__(self, base_url="http://10.0.6.170:9999", model="Qwen2.5-14B-Chat", embedding_model="bge-large-zh-v1.5") -> None:
        self.model = model
        self.embedding_model = embedding_model
        self.base_url = base_url
        self.headers = {'accept': 'application/json', 'Content-Type': 'application/json'}


    def chat(self, content):
        messages = [{"role": "user", "content": content}]
        data = {
            "model": self.model,
            "messages": messages,
            "temperature": 0.7,
            "top_p": 0.9,
            "stream": False
        }
        response = send_http(self.base_url, "/chat/completions", data, self.headers)
        return response['choices'][0]['message']['content']


    def embedings(self, content):
        messages = [content]
        data = {
            "model": self.embedding_model,
            "input": messages
        }
        response = send_http(self.base_url, "/embeddings", data, self.headers)
        return response['data'][0]['embedding']

